package com.aispeech.asr.clib.ssc;

import com.sun.jna.Native;
import lombok.extern.slf4j.Slf4j;

import java.util.List;

@Slf4j
public class Ssc {
    private static final String SSC_BIN_DEFAULT = "cn/bin/lasr_ssc_16000.bin";
    private static final String SSC_BIN_8K = "cn/bin/lasr_ssc_8000.bin";

    private static ISsc iSsc;
    private static long ptSscCfgDefault;
    private static long ptSscCfg8k;
    private long ptSscInstance;

    public Ssc() {
        this(16000);
    }

    public Ssc(int sampleRate) {
        long cfg = getscCfgSampleRate(sampleRate);
        ptSscInstance = iSsc.ssc_engine_new(cfg);
        if (ptSscInstance <= 0) {
            throw new RuntimeException("new SSC instance failed");
        }
    }

    public static void initialize(String resPath, String ipcPath) {
        iSsc = Native.load("ssc", ISsc.class);

        if (iSsc == null) {
            throw new RuntimeException("Load SSC library failed");
        }

        ptSscCfgDefault = iSsc.ssc_static_dependency_init(String.format("%s/%s", resPath, SSC_BIN_DEFAULT), ipcPath);
        if (ptSscCfgDefault <= 0) {
            throw new RuntimeException("Initialize SSC library failed");
        }

        ptSscCfg8k = iSsc.ssc_static_dependency_init(String.format("%s/%s", resPath, SSC_BIN_8K), ipcPath);
        if (ptSscCfg8k <= 0) {
            throw new RuntimeException("Initialize SSC library failed");
        }

        log.info("SSC library loaded.");
    }

    private static long getscCfgSampleRate(int sampleRate) {
        return (sampleRate == 8000) ? ptSscCfg8k : ptSscCfgDefault;
    }

    public int start(byte[] data, int nSpeakers, int nThreads, int total) {
        return iSsc.ssc_start(ptSscInstance, data, data.length, nSpeakers, nThreads, total);
    }

    public void feed(SscInput sscInput, int index) {
        List<SscInputElement> elements = sscInput.getElements();
        int size = elements.size();

        final SscInputStruct.ByReference sscInputStruct = new SscInputStruct.ByReference();
        sscInputStruct.aligns = new SscInputElementStruct.ByReference();
        SscInputElementStruct[] sscInputElementStructs = (SscInputElementStruct[]) sscInputStruct.aligns.toArray(size);
        sscInputStruct.len = size;

        for (int i = 0; i < size; i++) {
            SscInputElement element = elements.get(i);

            sscInputElementStructs[i].timeStart = element.getBegin();
            sscInputElementStructs[i].timeEnd = element.getEnd();
        }

        iSsc.ssc_feed(ptSscInstance, sscInputStruct, index);
    }

    public SscOutput stop() {
        SscOutputStruct.ByReference sscOutputStruct = iSsc.ssc_end(ptSscInstance);

        SscOutput output = new SscOutput(sscOutputStruct.errorCode, sscOutputStruct.procTime);

        if (output.getErrorCode() == 0) {
            SscOutputElementStruct[] sscOutputElementStructs = (SscOutputElementStruct[]) sscOutputStruct.sscResults.toArray(sscOutputStruct.len);
            for (SscOutputElementStruct outputElementStruct : sscOutputElementStructs) {
                output.addElement(outputElementStruct.timeStart, outputElementStruct.timeEnd, outputElementStruct.uid);
            }
        }

        return output;
    }

    @Override
    protected void finalize() throws Throwable {
        super.finalize();

        if (ptSscInstance != 0 && iSsc != null)
            iSsc.ssc_engine_delete(ptSscInstance);

        ptSscInstance = 0;
    }

    public static void destroy() {
        if (ptSscCfgDefault > 0)
            iSsc.ssc_static_dependency_destroy(ptSscCfgDefault);
        if (ptSscCfg8k > 0)
            iSsc.ssc_static_dependency_destroy(ptSscCfg8k);
    }
}
